import tensorflow as tf
import numpy as np
from math import sin, cos, radians, sqrt

# 初始化参数 (保持不变)
bit_num = 1
nRow = 64  # x方向单元个数，θ方向
mLine = 64  # y方向单元个数，φ方向
f0 = 300  # 工作频点，单位：GHz
c = 3e8  # 光速，单位：m/s
lamda0 = c / (f0 * 1e9)  # 波长，单位：m
dx = lamda0 / 2  # x方向单元间距，单位：m
dy = dx  # y方向单元间距，单位：m
k = 2 * np.pi / lamda0  # 波数

# 配置馈元和计算馈元位置
InF = 1  # 喇叭馈源摆放状态，0：垂直阵面，1：指向阵面中心
InTetha = 0  # 入射角
aphea_10dB = 45  # 喇叭馈源的-10dB波束角宽度

# 计算辐射方向图
qf = 1
qe = 1
eps = 0.000001



# 计算距离 (TensorFlow版本)
def distance_tf(p1, p2):
    return tf.sqrt(tf.reduce_sum(tf.square(p1 - p2)))


# 计算馈元 (保持不变)
def get_feed(InF, InTetha, aphea_10dB, nRow, dx):
    D = nRow * dx
    if InF == 0:
        beata = 90 - aphea_10dB / 2
        gema = 90 - beata - InTetha
        F = sin(radians(beata)) * ((D / 2) / sin(radians(gema)))
        xf = F * sin(radians(InTetha))
        yf = 0
        zf = F * cos(radians(InTetha))
    elif InF == 1:
        beata = 90 - InTetha
        gema = 180 - beata - aphea_10dB / 2
        F = sin(radians(gema)) * ((D / 2) / sin(radians(aphea_10dB / 2)))
        xf = F * sin(radians(InTetha))
        yf = 0
        zf = F * cos(radians(InTetha))
    return [xf, yf, zf], F / D



# 预计算馈元位置和F/D
Feed, FD = get_feed(InF, InTetha, aphea_10dB, nRow, dx)
Feed = tf.constant(Feed, dtype=tf.float32)

# 预计算天线位置坐标
posx = np.arange(-dx * (nRow - 1) / 2, dx * (nRow - 1) / 2 + dx, dx)
posy = np.arange(-dy * (mLine - 1) / 2, dy * (mLine - 1) / 2 + dy, dy)
posx = tf.constant(posx, dtype=tf.float32)
posy = tf.constant(posy, dtype=tf.float32)




# 计算方向图, fft方式, xyz坐标 (TensorFlow版本)
def get_pattern_fft_tf(phase, posx, posy, k, Feed, Rc, qf, qe):
    Nx, Ny = len(posx), len(posy)
    Kfft, Lfft = 360, 361

    # 生成相位矩阵 - 使用tf.complex来处理复数
    phase_matrix = tf.exp(tf.complex(0., phase))  # e^(i*phase) = cos(phase) + i*sin(phase)

    # 计算距离差分项
    distance_diff = []
    for i in range(Nx):
        for j in range(Ny):
            E = tf.stack([posx[i], posy[j], 0.0])
            rfij = distance_tf(Feed, E)
            distance_diff.append(-k * rfij)

    distance_diff = tf.reshape(tf.stack(distance_diff), (Nx, Ny))
    # 同样需要转换为复数
    phase_matrix = phase_matrix * tf.exp(tf.complex(0., distance_diff))

    # 使用FFT进行快速计算
    IF = tf.signal.fftshift(tf.signal.ifft2d(phase_matrix))
    F = IF * tf.cast(Kfft * Lfft, tf.complex64)  # 转换为复数类型

    # 返回绝对值作为方向图
    return tf.abs(F)


# 相位生成方向图, xy坐标系, FFT方式 (TensorFlow版本)
def phase_2_pattern_xyz_fft_tf(phaseBit):
    # 相位转角度再转弧度
    phaseBitRad = phaseBit * np.pi  # 因为180度=π弧度

    # 计算方向图
    pattern = get_pattern_fft_tf(phaseBitRad, posx, posy, k, Feed, [posx[0], posy[0], 0.0], qf, qe)
    return pattern


# TensorFlow可用的损失函数
def fitness_tf(y_true, y_pred):
    # 计算方向图
    true_pattern = phase_2_pattern_xyz_fft_tf(y_true)
    pred_pattern = phase_2_pattern_xyz_fft_tf(y_pred)

    # 计算均方误差损失
    loss = tf.reduce_mean(tf.square(true_pattern - pred_pattern))
    return loss


if __name__=="__main__":
    from util.util_log import setup_logging
    logger = setup_logging()
    # 创建测试数据
    y_true_tf = tf.random.normal([64, 64])
    y_pred_tf = tf.random.normal([64, 64])

    # 计算损失
    loss = fitness_tf(y_true_tf, y_pred_tf)
    logger.info(f"Loss: {loss.numpy()}")
    loss = fitness_tf(y_true_tf, y_true_tf)
    logger.info(f"Loss: {loss.numpy()}")

    # 梯度计算测试
    with tf.GradientTape() as tape:
        loss = fitness_tf(y_true_tf, y_pred_tf)
    grads = tape.gradient(loss, [y_pred_tf])

    print("Loss:", loss.numpy())
    print("Gradients:", grads)